#include "tcp_sender.hh"

#include "tcp_config.hh"

// #include <iostream>
#include <random>

// Dummy implementation of a TCP sender

// For Lab 3, please replace with a real implementation that passes the
// automated checks run by `make check_lab3`.

template <typename... Targs>
void DUMMY_CODE(Targs &&.../* unused */) {}

using namespace std;

//! \param[in] capacity the capacity of the outgoing byte stream
//! \param[in] retx_timeout the initial amount of time to wait before retransmitting the oldest outstanding segment
//! \param[in] fixed_isn the Initial Sequence Number to use, if set (otherwise uses a random ISN)
TCPSender::TCPSender(const size_t capacity, const uint16_t retx_timeout, const std::optional<WrappingInt32> fixed_isn)
    : _isn(fixed_isn.value_or(WrappingInt32{random_device()()}))
    , _initial_retransmission_timeout{retx_timeout}
    , _stream(capacity) {}

uint64_t TCPSender::bytes_in_flight() const { return _bytes_in_flight; }

void TCPSender::fill_window() {
    // if no window
    if (_window_size_other == 0)
        return;

    // if window shinked but ahead segment not sent out
    if (_next_seqno >= (_window_idx + _window_size_other))
        return;

    // something has sent
    bool _is_empty_before = _bytes_in_flight == 0;
    // real window size
    uint64_t real_window_size = _window_is_zero ? 1 : _window_size_other;
    // if SYN segment is not sent
    if (__builtin_expect(!_is_SYN_sent, false)) {
        // _next_seqno != 0 || _bytes_in_flight != 0 || real_ws != 1
        // throw Exception
        if (real_window_size != 1)
            throw std::logic_error("window size doesn't match initial state");
        if (_bytes_in_flight != 0)
            throw std::logic_error("in flight byte num doesn't match initial state");
        if (_next_seqno != 0)
            throw std::logic_error("next seqno doesn't match initial state");
        _is_SYN_sent = true;
        TCPSegment syn_seg;

        syn_seg.header().syn = true;
        syn_seg.header().seqno = _isn;

        _bytes_in_flight += syn_seg.header().syn;
        _next_seqno += syn_seg.header().syn;

        _segments_out.push(syn_seg);
        _segments_trace.emplace(syn_seg, _timer);
    }

    else {
        // calculate the remain size in window
        uint64_t window_remain = real_window_size > _bytes_in_flight ? real_window_size - _bytes_in_flight : 0;

        while (window_remain && !_is_FIN_sent) {
            uint64_t expect_fin_segment_payload_size = _stream.buffer_size();
            bool can_generate_fin = _stream.input_ended() && (!_is_FIN_sent) && window_remain >= 1 &&
                                    expect_fin_segment_payload_size <= TCPConfig::MAX_PAYLOAD_SIZE &&
                                    window_remain >= (expect_fin_segment_payload_size + 1);
            if (__builtin_expect(can_generate_fin, false)) {
                // condition to sent FIN sengment :
                // 1. stream is input_ended
                // 2. stream remain size <= MAX_PAYLOAD (for FIN flag)
                // 3. window remain now  >= (stream remain size + 1) (for fin flag)
                // 4. the sender not sent FIN before
                // 5. window has space
                _is_FIN_sent = true;

                TCPSegment fin_seg;

                fin_seg.header().fin = true;
                fin_seg.header().seqno = wrap(_next_seqno, _isn);

                fin_seg.payload() = Buffer(_stream.read(expect_fin_segment_payload_size));

                _bytes_in_flight += expect_fin_segment_payload_size + 1;
                _next_seqno += expect_fin_segment_payload_size + 1;

                _segments_out.push(fin_seg);
                _segments_trace.emplace(fin_seg, _timer);

                break;
            } else {
                /* Construct the segment */
                // construct payload
                uint64_t read_size = std::min({window_remain, TCPConfig::MAX_PAYLOAD_SIZE, _stream.buffer_size()});
                // no payload , exit
                if (read_size == 0)
                    break;
                std::string payload = _stream.read(read_size);
                // initialize segment
                TCPSegment newSeg;
                newSeg.payload() = Buffer(std::move(payload));
                newSeg.header().seqno = wrap(_next_seqno, _isn);

                // send
                _segments_out.push(newSeg);
                _segments_trace.emplace(newSeg, _timer);

                // update internal state
                _bytes_in_flight += read_size;
                _next_seqno += read_size;

                // update window remain
                window_remain -= read_size;
            }
        }
    }
    if (_is_empty_before && _bytes_in_flight > 0) {
        _retransmission_timeout = _initial_retransmission_timeout;
        _retransmission_times = 0;
    }
    return;
}

//! \param ackno The remote receiver's ackno (acknowledgment number)
//! \param window_size The remote receiver's advertised window size
void TCPSender::ack_received(const WrappingInt32 ackno, const uint16_t window_size) {
    // calculate the absolute idx of the ack
    uint64_t abs_ack_no = unwrap(ackno, _isn, _window_idx);
    // if the abs_ack_no is behind of next_seqno
    // discard for unreasonable
    if (abs_ack_no > _next_seqno)
        return;

    if (!_segments_trace.empty()) {
        // get out the information of first segment in flight trace queue
        // to initialize the loop parameter
        WrappingInt32 seqno = _segments_trace.front().seg.header().seqno;
        uint64_t segment_size = _segments_trace.front().seg.payload().size();
        bool syn = _segments_trace.front().seg.header().syn;
        bool fin = _segments_trace.front().seg.header().fin;

        uint64_t abs_seq_top = unwrap(seqno, _isn, _window_idx);
        uint64_t fin_or_syn = fin + syn;

        // accumulated ack, pop all segment whose seqno
        // less than abs_ack_no
        while (abs_seq_top < abs_ack_no) {
            // not fully acknowledged , stop loop
            if (abs_ack_no < (abs_seq_top + segment_size + fin_or_syn))
                break;
            else {
                // update the window state parameter
                _bytes_in_flight -= (segment_size + fin_or_syn);
                _window_idx += (segment_size + fin_or_syn);
                // pop
                _segments_trace.pop();
                // set timer of the new flight queue
                if (!_segments_trace.empty())
                    _segments_trace.front().sent_time = _timer;
                _retransmission_times = 0;
                _retransmission_timeout = _initial_retransmission_timeout;
            }

            // flight queue empty , stop loop
            if (_segments_trace.empty()) {
                break;
            }

            // update loop parameters
            syn = _segments_trace.front().seg.header().syn;
            fin = _segments_trace.front().seg.header().fin;

            seqno = _segments_trace.front().seg.header().seqno;
            segment_size = _segments_trace.front().seg.payload().size();

            abs_seq_top = unwrap(seqno, _isn, _window_idx);
            fin_or_syn = fin + syn;
        }
    }
    // update window size
    _window_is_zero = window_size == 0;
    _window_size_other = window_size > 0 ? window_size : 1;
}

//! \param[in] ms_since_last_tick the number of milliseconds since the last call to this method
void TCPSender::tick(const size_t ms_since_last_tick) {
    // update time by pluse ms_since_last_tick
    _timer += ms_since_last_tick;

    // a timeout timer is enabled
    if (_bytes_in_flight > 0) {
        // see whether the first segment in the flight queue
        // is expired
        uint64_t timeout_start = _segments_trace.front().sent_time;
        if (_timer - timeout_start >= _retransmission_timeout) {
            // expired , then retransmit
            _segments_out.push(_segments_trace.front().seg);
            // whether the window size is zero since last ack
            if (!_window_is_zero) {
                _retransmission_times += 1;
                _retransmission_timeout = _initial_retransmission_timeout * (1 << (_retransmission_times));
            }
            _segments_trace.front().sent_time = _timer;
        }
    }
    return;
}

unsigned int TCPSender::consecutive_retransmissions() const { return _retransmission_times; }

void TCPSender::send_empty_segment() {
    TCPSegment newSeg;
    newSeg.header().seqno = wrap(_next_seqno, _isn);
    newSeg.payload() = Buffer("");

    _segments_out.push(newSeg);
}
